import numpy as np
"""
basic matrix multiplication
"""
class MetaDense:
    def __init__(self, input_channel, output_channel, init_str=None, init_lambda=0):
        self.w = 2*np.random.random((input_channel, output_channel)) - 1
        if init_str == "Xavier":
            self.w = np.random.randn(input_channel, output_channel) / np.sqrt(input_channel)
        elif init_str == "He":
            self.w = np.random.randn(input_channel, output_channel) * np.sqrt(2 / input_channel)
        self.x = None
        self.dw = None
        self.dx = None
        self.init_lambda = init_lambda

    def __call__(self, x):
        self.x = x
        return self.x @ self.w

    def backward(self, dout):
        self.dx = dout @ self.w.transpose()
        self.dw = self.x.reshape(-1, 1) @ dout + self.init_lambda * self.w
        return self.dx, self.dw